/*
*  Copyright (c) 2016 The WebM project authors. All Rights Reserved.
*
*  Use of this source code is governed by a BSD-style license
*  that can be found in the LICENSE file in the root of the source
*  tree. An additional intellectual property rights grant can be found
*  in the file PATENTS.  All contributing project authors may
*  be found in the AUTHORS file in the root of the source tree.
*/

#include <math.h>
#include <assert.h>
//#include <vp8/common/blockd.h>
#include <vp9/common/vp9_enums.h>
#include <vp9/common/vp9_blockd.h>
#include "./vpx_dsp_rtcd.h"
#include "vpx_dsp/psnr.h"
#include "vpx_scale/yv12config.h"

double vpx_sse_to_psnr(double samples, double peak, double sse) {
    if (sse > 0.0) {
        const double psnr = 10.0 * log10(samples * peak * peak / sse);
        return psnr > MAX_PSNR ? MAX_PSNR : psnr;
    } else {
        return MAX_PSNR;
    }
}

/* TODO(yaowu): The block_variance calls the unoptimized versions of variance()
* and highbd_8_variance(). It should not.
*/
static void encoder_variance(const uint8_t *a, int a_stride, const uint8_t *b,
                             int b_stride, int w, int h, unsigned int *sse,
                             int *sum) {
    int i, j;

    *sum = 0;
    *sse = 0;

    for (i = 0; i < h; i++) {
        for (j = 0; j < w; j++) {
            const int diff = a[j] - b[j];
            *sum += diff;
            *sse += diff * diff;
        }

        a += a_stride;
        b += b_stride;
    }
}

#if CONFIG_VP9_HIGHBITDEPTH
static void encoder_highbd_variance64(const uint8_t *a8, int a_stride,
                                      const uint8_t *b8, int b_stride, int w,
                                      int h, uint64_t *sse, int64_t *sum) {
  int i, j;

  uint16_t *a = CONVERT_TO_SHORTPTR(a8);
  uint16_t *b = CONVERT_TO_SHORTPTR(b8);
  *sum = 0;
  *sse = 0;

  for (i = 0; i < h; i++) {
    for (j = 0; j < w; j++) {
      const int diff = a[j] - b[j];
      *sum += diff;
      *sse += diff * diff;
    }
    a += a_stride;
    b += b_stride;
  }
}

static void encoder_highbd_8_variance(const uint8_t *a8, int a_stride,
                                      const uint8_t *b8, int b_stride, int w,
                                      int h, unsigned int *sse, int *sum) {
  uint64_t sse_long = 0;
  int64_t sum_long = 0;
  encoder_highbd_variance64(a8, a_stride, b8, b_stride, w, h, &sse_long,
                            &sum_long);
  *sse = (unsigned int)sse_long;
  *sum = (int)sum_long;
}
#endif  // CONFIG_VP9_HIGHBITDEPTH

static int64_t get_sse(const uint8_t *a, int a_stride, const uint8_t *b,
                       int b_stride, int width, int height) {
    const int dw = width % 16;
    const int dh = height % 16;
    int64_t total_sse = 0;
    unsigned int sse = 0;
    int sum = 0;
    int x, y;

    if (dw > 0) {
        encoder_variance(&a[width - dw], a_stride, &b[width - dw], b_stride, dw,
                         height, &sse, &sum);
        total_sse += sse;
    }

    if (dh > 0) {
        encoder_variance(&a[(height - dh) * a_stride], a_stride,
                         &b[(height - dh) * b_stride], b_stride, width - dw, dh,
                         &sse, &sum);
        total_sse += sse;
    }

    for (y = 0; y < height / 16; ++y) {
        const uint8_t *pa = a;
        const uint8_t *pb = b;
        for (x = 0; x < width / 16; ++x) {
            vpx_mse16x16(pa, a_stride, pb, b_stride, &sse);
            total_sse += sse;

            pa += 16;
            pb += 16;
        }

        a += 16 * a_stride;
        b += 16 * b_stride;
    }

    return total_sse;
}

#if CONFIG_VP9_HIGHBITDEPTH
static int64_t highbd_get_sse_shift(const uint8_t *a8, int a_stride,
                                    const uint8_t *b8, int b_stride, int width,
                                    int height, unsigned int input_shift) {
  const uint16_t *a = CONVERT_TO_SHORTPTR(a8);
  const uint16_t *b = CONVERT_TO_SHORTPTR(b8);
  int64_t total_sse = 0;
  int x, y;
  for (y = 0; y < height; ++y) {
    for (x = 0; x < width; ++x) {
      int64_t diff;
      diff = (a[x] >> input_shift) - (b[x] >> input_shift);
      total_sse += diff * diff;
    }
    a += a_stride;
    b += b_stride;
  }
  return total_sse;
}

static int64_t highbd_get_sse(const uint8_t *a, int a_stride, const uint8_t *b,
                              int b_stride, int width, int height) {
  int64_t total_sse = 0;
  int x, y;
  const int dw = width % 16;
  const int dh = height % 16;
  unsigned int sse = 0;
  int sum = 0;
  if (dw > 0) {
    encoder_highbd_8_variance(&a[width - dw], a_stride, &b[width - dw],
                              b_stride, dw, height, &sse, &sum);
    total_sse += sse;
  }
  if (dh > 0) {
    encoder_highbd_8_variance(&a[(height - dh) * a_stride], a_stride,
                              &b[(height - dh) * b_stride], b_stride,
                              width - dw, dh, &sse, &sum);
    total_sse += sse;
  }
  for (y = 0; y < height / 16; ++y) {
    const uint8_t *pa = a;
    const uint8_t *pb = b;
    for (x = 0; x < width / 16; ++x) {
      vpx_highbd_8_mse16x16(pa, a_stride, pb, b_stride, &sse);
      total_sse += sse;
      pa += 16;
      pb += 16;
    }
    a += 16 * a_stride;
    b += 16 * b_stride;
  }
  return total_sse;
}
#endif  // CONFIG_VP9_HIGHBITDEPTH

int64_t vpx_get_y_sse(const YV12_BUFFER_CONFIG *a,
                      const YV12_BUFFER_CONFIG *b) {
    assert(a->y_crop_width == b->y_crop_width);
    assert(a->y_crop_height == b->y_crop_height);

    return get_sse(a->y_buffer, a->y_stride, b->y_buffer, b->y_stride,
                   a->y_crop_width, a->y_crop_height);
}

#if CONFIG_VP9_HIGHBITDEPTH
int64_t vpx_highbd_get_y_sse(const YV12_BUFFER_CONFIG *a,
                             const YV12_BUFFER_CONFIG *b) {
  assert(a->y_crop_width == b->y_crop_width);
  assert(a->y_crop_height == b->y_crop_height);
  assert((a->flags & YV12_FLAG_HIGHBITDEPTH) != 0);
  assert((b->flags & YV12_FLAG_HIGHBITDEPTH) != 0);

  return highbd_get_sse(a->y_buffer, a->y_stride, b->y_buffer, b->y_stride,
                        a->y_crop_width, a->y_crop_height);
}
#endif  // CONFIG_VP9_HIGHBITDEPTH

#if CONFIG_VP9_HIGHBITDEPTH
void vpx_calc_highbd_psnr(const YV12_BUFFER_CONFIG *a,
                          const YV12_BUFFER_CONFIG *b, PSNR_STATS *psnr,
                          uint32_t bit_depth, uint32_t in_bit_depth) {
  const int widths[3] = { a->y_crop_width, a->uv_crop_width, a->uv_crop_width };
  const int heights[3] = { a->y_crop_height, a->uv_crop_height,
                           a->uv_crop_height };
  const uint8_t *a_planes[3] = { a->y_buffer, a->u_buffer, a->v_buffer };
  const int a_strides[3] = { a->y_stride, a->uv_stride, a->uv_stride };
  const uint8_t *b_planes[3] = { b->y_buffer, b->u_buffer, b->v_buffer };
  const int b_strides[3] = { b->y_stride, b->uv_stride, b->uv_stride };
  int i;
  uint64_t total_sse = 0;
  uint32_t total_samples = 0;
  const double peak = (double)((1 << in_bit_depth) - 1);
  const unsigned int input_shift = bit_depth - in_bit_depth;

  for (i = 0; i < 3; ++i) {
    const int w = widths[i];
    const int h = heights[i];
    const uint32_t samples = w * h;
    uint64_t sse;
    if (a->flags & YV12_FLAG_HIGHBITDEPTH) {
      if (input_shift) {
        sse = highbd_get_sse_shift(a_planes[i], a_strides[i], b_planes[i],
                                   b_strides[i], w, h, input_shift);
      } else {
        sse = highbd_get_sse(a_planes[i], a_strides[i], b_planes[i],
                             b_strides[i], w, h);
      }
    } else {
      sse = get_sse(a_planes[i], a_strides[i], b_planes[i], b_strides[i], w, h);
    }
    psnr->sse[1 + i] = sse;
    psnr->samples[1 + i] = samples;
    psnr->psnr[1 + i] = vpx_sse_to_psnr(samples, peak, (double)sse);

    total_sse += sse;
    total_samples += samples;
  }

  psnr->sse[0] = total_sse;
  psnr->samples[0] = total_samples;
  psnr->psnr[0] =
      vpx_sse_to_psnr((double)total_samples, peak, (double)total_sse);
}

#endif  // !CONFIG_VP9_HIGHBITDEPTH

void vpx_calc_psnr(const YV12_BUFFER_CONFIG *a, const YV12_BUFFER_CONFIG *b,
                   PSNR_STATS *psnr) {
    static const double peak = 255.0;
    const int widths[3] = {a->y_crop_width, a->uv_crop_width, a->uv_crop_width};
    const int heights[3] = {a->y_crop_height, a->uv_crop_height,
                            a->uv_crop_height};
    const uint8_t *a_planes[3] = {a->y_buffer, a->u_buffer, a->v_buffer};
    const int a_strides[3] = {a->y_stride, a->uv_stride, a->uv_stride};
    const uint8_t *b_planes[3] = {b->y_buffer, b->u_buffer, b->v_buffer};
    const int b_strides[3] = {b->y_stride, b->uv_stride, b->uv_stride};
    int i;
    uint64_t total_sse = 0;
    uint32_t total_samples = 0;

    for (i = 0; i < 3; ++i) {
        const int w = widths[i];
        const int h = heights[i];
        const uint32_t samples = w * h;
        const uint64_t sse =
                get_sse(a_planes[i], a_strides[i], b_planes[i], b_strides[i], w, h);
        psnr->sse[1 + i] = sse;
        psnr->samples[1 + i] = samples;
        psnr->psnr[1 + i] = vpx_sse_to_psnr(samples, peak, (double) sse);

        total_sse += sse;
        total_samples += samples;
    }

    psnr->sse[0] = total_sse;
    psnr->samples[0] = total_samples;
    psnr->psnr[0] =
            vpx_sse_to_psnr((double) total_samples, peak, (double) total_sse);
}

//void vpx_calc_psnr_block_old(const YV12_BUFFER_CONFIG *a, const YV12_BUFFER_CONFIG *b,
//                        PSNR_STATS *psnr, int width, int height, int mi_col, int mi_row) {
//    static const double peak = 255.0;
//    const int x0 = mi_col * MI_SIZE;
//    const int y0 = mi_row * MI_SIZE;
//
//    const int x_offsets[3] = {x0, x0 >> a->subsampling_x, x0 >> a->subsampling_x};
//    const int y_offsets[3] = {y0, y0 >> a->subsampling_y, y0 >> a->subsampling_y};
//    const int widths[3] = {width, width >> a->subsampling_x, width >> a->subsampling_x};
//    const int heights[3] = {height, height >> a->subsampling_y, height >> a->subsampling_x};
//    const int a_strides[3] = {a->y_stride, a->uv_stride, a->uv_stride};
//    const uint8_t *a_planes[3] = {&a->y_buffer[y_offsets[0] * a_strides[0] + x_offsets[0]],
//                                  &a->u_buffer[y_offsets[1] * a_strides[1] + x_offsets[1]],
//                                  &a->v_buffer[y_offsets[2] * a_strides[2] + x_offsets[2]],
//    };
//    const int b_strides[3] = {b->y_stride, b->uv_stride, b->uv_stride};
//    const uint8_t *b_planes[3] = {&b->y_buffer[y_offsets[0] * b_strides[0] + x_offsets[0]],
//                                  &b->u_buffer[y_offsets[1] * b_strides[1] + x_offsets[1]],
//                                  &b->v_buffer[y_offsets[2] * b_strides[2] + x_offsets[2]],
//    };
//
////  int i, x, y;
////  int error;
////  int total_mse = 0;
////  for (i = 0; i < 3; i++) {
////      for (x = 0; x < widths[i]; x++) {
////          for (y = 0; y < heights[i]; y ++) {
////              error = a_planes[i][y * a_strides[i] + x] - b_planes[i][y * b_strides[i] + x];
////              total_mse += pow(error, 2);
////          }
////      }
////  }
//
////  return total_mse;
//
//
//    int i;
//    uint64_t total_sse = 0;
//    uint32_t total_samples = 0;
//
//    for (i = 0; i < 3; ++i) {
//        const int w = widths[i];
//        const int h = heights[i];
//        const uint32_t samples = w * h;
//        const uint64_t sse =
//                get_sse(a_planes[i], a_strides[i], b_planes[i], b_strides[i], w, h);
//        psnr->sse[1 + i] = sse;
//        psnr->samples[1 + i] = samples;
//        psnr->psnr[1 + i] = vpx_sse_to_psnr(samples, peak, (double) sse);
//
//        total_sse += sse;
//        total_samples += samples;
//    }
//
//    psnr->sse[0] = total_sse;
//    psnr->samples[0] = total_samples;
//    psnr->psnr[0] =
//            vpx_sse_to_psnr((double) total_samples, peak, (double) total_sse);
//}

void vpx_calc_block_psnr(const YV12_BUFFER_CONFIG *a, const YV12_BUFFER_CONFIG *b,
                         PSNR_STATS *psnr, const int *widths, const int *heights, const int *x_offsets, const int *y_offsets, int scale) {
    static const double peak = 255.0;

    const int a_strides[3] = {a->y_stride, a->uv_stride, a->uv_stride};
    const uint8_t *a_planes[3] = {&a->y_buffer[y_offsets[0] * a_strides[0] * scale + x_offsets[0] * scale],
                                  &a->u_buffer[y_offsets[1] * a_strides[1] * scale + x_offsets[1] * scale],
                                  &a->v_buffer[y_offsets[2] * a_strides[2] * scale + x_offsets[2] * scale],
    };
    const int b_strides[3] = {b->y_stride, b->uv_stride, b->uv_stride};
    const uint8_t *b_planes[3] = {&b->y_buffer[y_offsets[0] * b_strides[0] * scale + x_offsets[0] * scale],
                                  &b->u_buffer[y_offsets[1] * b_strides[1] * scale + x_offsets[1] * scale],
                                  &b->v_buffer[y_offsets[2] * b_strides[2] * scale + x_offsets[2] * scale],
    };

    int i;
    uint64_t total_sse = 0;
    uint32_t total_samples = 0;

    for (i = 0; i < 3; ++i) {
        const int w = widths[i] * scale;
        const int h = heights[i] * scale;
        const uint32_t samples = w * h;
        const uint64_t sse =
                get_sse(a_planes[i], a_strides[i], b_planes[i], b_strides[i], w, h);
        psnr->sse[1 + i] = sse;
        psnr->samples[1 + i] = samples;
        psnr->psnr[1 + i] = vpx_sse_to_psnr(samples, peak, (double) sse);

        total_sse += sse;
        total_samples += samples;
    }

    psnr->sse[0] = total_sse;
    psnr->samples[0] = total_samples;
    psnr->psnr[0] =
            vpx_sse_to_psnr((double) total_samples, peak, (double) total_sse);
}
